/*
 * Copyright (C) 2023 Coder.AN
 * Email: an.hongjun@foxmail.com
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
#include "edge_tensor.h"
#include "edge_tensor_kernel.h"
#include <stdio.h>
#include <math.h>
#include <thrust/extrema.h>


namespace draw{
__device__ void cal_dim(int idx, int* n, int* h, int* w, int*c, Dim dim, DimOrder order)
{
    if (order == NHWC)
    {
        *n = idx / (dim.H * dim.W * dim.C);
        idx -= *n * dim.H * dim.W * dim.C;
        *h = idx / (dim.W * dim.C);
        idx -= *h * dim.W * dim.C;
        *w = idx / dim.C;
        idx -= *w * dim.C;
        *c = idx;
    }
    else if (order == NCWH)
    {
        *n = idx / (dim.C * dim.W * dim.H);
        idx -= *n * dim.C * dim.W * dim.H;
        *c = idx / (dim.W * dim.H);
        idx -= *c * dim.W * dim.H;
        *w = idx / dim.H;
        idx -= *w * dim.H;
        *h = idx;
    }
    else if (order == NCHW)
    {
        *n = idx / (dim.C * dim.H * dim.W);
        idx -= *n * dim.C * dim.H * dim.W;
        *c = idx / (dim.H * dim.W);
        idx -= *c * dim.H * dim.W;
        *h = idx / dim.W;
        idx -= *h * dim.W;
        *w = idx;
    }
}

__device__ int cal_new_idx(int n, int h, int w, int c, Dim dim, DimOrder order)
{
    if (order == NHWC)
    {
        return n * dim.H * dim.W * dim.C + h * dim.W * dim.C + w * dim.C + c;
    }
    else if (order == NCWH)
    {
        return n * dim.C * dim.W * dim.H + c * dim.W * dim.H + w * dim.H + h;
    }
    else if (order == NCHW)
    {
        return n * dim.C * dim.H * dim.W + c * dim.H * dim.W + h * dim.W + w;
    }
    return 0;
}
}

__global__ void DrawGrayLine_kernel(uint8_t* dst, size_t unit_count, Dim dim, DimOrder dim_order, 
                                    uint8_t grayVal, int x_min, int x_max, int y_min, int y_max, int A, int B, int C, int r, float b)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= unit_count) return;

    int n, h, w, c;
    draw::cal_dim(idx, &n, &h, &w, &c, dim, dim_order);

    float a = A * h + B * w + C;
    if (a < 0) a = -a;

    float d = a / b;
    bool w_mask = (h >= x_min && h <= x_max);
    bool h_mask = (w >= y_min && w <= y_max);
    bool mask = (w_mask && h_mask && d <= r + 1);

    if (mask)
    {
        dst[idx] = grayVal;
    }
}

void DrawGrayLine(uint8_t* dst, size_t unit_count, Dim dim, DimOrder dim_order, 
                    uint8_t grayVal, int x_min, int x_max, int y_min, int y_max, int A, int B, int C, int r)
{
    dim3 block(1024);
    dim3 grid((unit_count - 1) / block.x + 1);
    float b = sqrt(A * A + B * B);
    DrawGrayLine_kernel<<<grid, block>>>(dst, unit_count, dim, dim_order, grayVal, 
                                            x_min, x_max, y_min, y_max, A, B, C, r, b);
    CHECK(cudaDeviceSynchronize());
    return;
}

__global__ void DrawColorLine_kernel(uint8_t* dst, size_t unit_count, Dim dim, DimOrder dim_order, 
                                        Color colorVal, int x_min, int x_max, int y_min, int y_max, int A, int B, int C, int r, float b)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= unit_count) return;

    int n, h, w, c;
    draw::cal_dim(idx, &n, &h, &w, &c, dim, dim_order);

    float a = A * h + B * w + C;
    if (a < 0) a = -a;
    
    float d = a / b;
    bool w_mask = (h >= x_min && h <= x_max);
    bool h_mask = (w >= y_min && w <= y_max);
    bool mask = (w_mask && h_mask && d <= r + 1);

    if (mask)
    {
        dst[idx] = ((uint8_t*)&colorVal)[c];
    }
}

void DrawColorLine(uint8_t* dst, size_t unit_count, Dim dim, DimOrder dim_order, 
                    Color colorVal, int x_min, int x_max, int y_min, int y_max, int A, int B, int C, int r)
{
    dim3 block(1024);
    dim3 grid((unit_count - 1) / block.x + 1);
    float b = sqrt(A * A + B * B);
    DrawColorLine_kernel<<<grid, block>>>(dst, unit_count, dim, dim_order, colorVal, 
                                            x_min, x_max, y_min, y_max, A, B, C, r, b);
    CHECK(cudaDeviceSynchronize());
    return;
}

__global__ void DrawGrayCircle_kernel(uint8_t* dst, size_t unit_count, Dim dim, DimOrder dim_order, 
                                        Point center, uint8_t grayVal, int r2)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= unit_count) return;

    int n, h, w, c;
    draw::cal_dim(idx, &n, &h, &w, &c, dim, dim_order);

    float d = (center.h - h) * (center.h - h);
    d += (center.w - w) * (center.w - w);

    if (d <= r2+1)
    {
        dst[idx] = grayVal;
    }
}

void DrawGrayCircle(uint8_t* dst, size_t unit_count, Dim dim, DimOrder dim_order,
                        Point center, uint8_t grayVal, int r2)
{
    dim3 block(1024);
    dim3 grid((unit_count - 1) / block.x + 1);
    DrawGrayCircle_kernel<<<grid, block>>>(dst, unit_count, dim, dim_order, 
                                            center, grayVal, r2);
    CHECK(cudaDeviceSynchronize());
    return;
}

__global__ void DrawColorCircle_kernel(uint8_t* dst, size_t unit_count, Dim dim, DimOrder dim_order, 
                                        Point center, Color colorVal, int r2)
{
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= unit_count) return;

    int n, h, w, c;
    draw::cal_dim(idx, &n, &h, &w, &c, dim, dim_order);

    float d = (center.h - h) * (center.h - h);
    d += (center.w - w) * (center.w - w);

    if (d <= r2+1)
    {
        dst[idx] = ((uint8_t*)&colorVal)[c];
    }
}

void DrawColorCircle(uint8_t* dst, size_t unit_count, Dim dim, DimOrder dim_order,
                        Point center, Color colorVal, int r2)
{
    dim3 block(1024);
    dim3 grid((unit_count - 1) / block.x + 1);
    DrawColorCircle_kernel<<<grid, block>>>(dst, unit_count, dim, dim_order, 
                                            center, colorVal, r2);
    CHECK(cudaDeviceSynchronize());
    return;
}
